
import os
import torch
from argparse import ArgumentParser, Namespace

from yaml import parse

def get_args_sbi():
   parser = ArgumentParser()
   parser.add_argument('--lr', '-l', help='learning rate', default=1e-3)
   parser.add_argument('--batch_size', '-b', default=2, type=int)
   parser.add_argument('--steps_for_eval', default=1, type=int)
   parser.add_argument('--steps_eval', default=1, type=int)
   parser.add_argument('--epochs', default=10, type=int)
   parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
   parser.add_argument('--gpu_ids', default='0', type=str)
   
   parser.add_argument('--resume', default=False, type=bool)
   parser.add_argument('--resume_epoch', default=9)
   parser.add_argument('--mode', default='train')
   parser.add_argument('--save_error', default=False, type=bool)
   parser.add_argument('--save_predict', default=False, type=bool)
   
   # parser.add_argument('--data_dir', default='E:/ll/data/ps/split_train_test')
   # parser.add_argument('--data_dir', default='F:/Sources/datasets/face/split_train_test') # 'F:/VOS/ubutu-vm/resources/datasets/work/train') 
   parser.add_argument('--data_dir', default='') # 'F:/VOS/ubutu-vm/resources/datasets/work/train') 
   parser.add_argument('--train_dir', default='/home/algtest/datasets/ps_dataset/train')
   parser.add_argument('--test_dir', default='/home/algtest/datasets/ps_dataset/test')
   parser.add_argument('--landmarks_path', default='/home/algtest/datasets/ps_dataset/landmarks.json')
   parser.add_argument('--train_faces_path', default='/home/algtest/datasets/ps_dataset/train_faces.json')
   parser.add_argument('--test_faces_path', default='/home/algtest/datasets/ps_dataset/test_faces.json')
   parser.add_argument('--img_size', default='240, 240')
   parser.add_argument('--crop_scale', default='0.9, 1')
   parser.add_argument('--writer_type', default='tensorboard')
   
   # 一些路徑
   parser.add_argument('--log_dir', default='./log')
   parser.add_argument('--model_dir', default='./model')
   parser.add_argument('--predict_dir', default='./predict')
   parser.add_argument('--error_dir', default='./error')
   
   
   parser.add_argument('--log_file', default='run.log')
   args = parser.parse_args()
   
   # 處理下 args
   args.img_size = [int(v) for v in args.img_size.split(',')]
   args.crop_scale = [float(v) for v in args.crop_scale.split(',')]
   for k, v in args._get_kwargs():
      if v and k.endswith('_dir') and not os.path.exists(v):
         os.makedirs(v)
   return args



def get_args():
   parser = ArgumentParser()
   parser.add_argument('--lr', '-l', help='learning rate', default=1e-3)
   parser.add_argument('--batch_size', '-b', default=2)
   parser.add_argument('--steps_for_eval', default=1)
   parser.add_argument('--steps_eval', default=1)
   parser.add_argument('--epochs', default=10)
   parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
   
   parser.add_argument('--resume', default=False)
   parser.add_argument('--resume_epoch', default=9)
   
   # parser.add_argument('--data_dir', default='E:/ll/data/ps/split_train_test')
   # parser.add_argument('--data_dir', default='F:/Sources/datasets/face/split_train_test') # 'F:/VOS/ubutu-vm/resources/datasets/work/train') 
   parser.add_argument('--data_dir', default='F:/VOS/ubutu-vm/resources/datasets/work/split_train_test') # 'F:/VOS/ubutu-vm/resources/datasets/work/train') 
   parser.add_argument('--train_dir', default=['gen_train', 'train'])
   parser.add_argument('--test_dir', default=['gen_test', 'test'])
   parser.add_argument('--log_dir', default='./log')
   parser.add_argument('--log_file', default='run.log')
   parser.add_argument('--writer_type', default='tensorboard')
   return parser.parse_args()